from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.remote.webdriver import WebDriver
from typing import Optional, Dict, List
import time
import logging
from enum import Enum
import os

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class BrowserType(Enum):
    CHROME = "chrome"
    FIREFOX = "firefox"
    EDGE = "edge"
    SAFARI = "safari"

class DriverManager:
    """浏览器驱动管理器，负责创建、配置和管理WebDriver实例"""
    def __init__(self):
        self.active_drivers: Dict[str, WebDriver] = {}
        self.driver_counter = 0
        
    def create_driver(self, browser_type: BrowserType = BrowserType.CHROME, 
                     headless: bool = False, 
                     window_size: tuple = (1920, 1080),
                     user_agent: Optional[str] = None,
                     extra_options: Optional[List[str]] = None) -> str:
        """
        创建新的浏览器驱动实例
        :param browser_type: 浏览器类型
        :param headless: 是否无头模式
        :param window_size: 窗口大小 (宽, 高)
        :param user_agent: 自定义User-Agent
        :param extra_options: 额外的浏览器选项
        :return: 驱动ID
        """
        driver_id = f"driver_{self.driver_counter}"
        self.driver_counter += 1
        
        try:
            if browser_type == BrowserType.CHROME:
                options = Options()
                
                # 基础配置
                if headless:
                    options.add_argument("--headless=new")
                options.add_argument(f"--window-size={window_size[0]},{window_size[1]}")
                
                # 性能优化配置
                options.add_argument("--disable-gpu")
                options.add_argument("--no-sandbox")
                options.add_argument("--disable-dev-shm-usage")
                
                # 自定义User-Agent
                if user_agent:
                    options.add_argument(f"--user-agent={user_agent}")
                
                # 额外选项
                if extra_options:
                    for option in extra_options:
                        options.add_argument(option)
                
                # 创建驱动，使用相对路径
                current_dir = os.path.dirname(os.path.abspath(__file__))
                chromedriver_path = os.path.join(current_dir, '..', 'chromedriver.exe')
                
                # 增强Chrome浏览器二进制文件路径检测逻辑
                chrome_binary_paths = [
                    'C:/Program Files/Google/Chrome/Application/chrome.exe',
                    'C:/Program Files (x86)/Google/Chrome/Application/chrome.exe',
                    f'{os.path.expanduser("~")}/AppData/Local/Google/Chrome/Application/chrome.exe',
                    'C:/Program Files/Google/Chrome Dev/Application/chrome.exe',
                    'C:/Program Files/Google/Chrome Beta/Application/chrome.exe',
                    'C:/Program Files/Google/Chrome Canary/Application/chrome.exe',
                    'C:/Windows/chrome.exe'
                ]
                
                # 尝试检测Chrome路径
                chrome_found = False
                for path in chrome_binary_paths:
                    if os.path.exists(path):
                        options.binary_location = path
                        chrome_found = True
                        logger.info(f"找到Chrome二进制文件: {path}")
                        break
                
                if not chrome_found:
                    logger.warning("未找到Chrome二进制文件，将尝试使用默认路径或系统PATH中的Chrome")
                    # 尝试自动查找Chrome安装路径（通过注册表或其他方式）
                    try:
                        import winreg
                        reg_path = r'SOFTWARE\Microsoft\Windows\CurrentVersion\App Paths\chrome.exe'
                        with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, reg_path) as key:
                            chrome_path = winreg.QueryValue(key, None)
                            if chrome_path and os.path.exists(chrome_path):
                                options.binary_location = chrome_path
                                logger.info(f"通过注册表找到Chrome: {chrome_path}")
                                chrome_found = True
                    except Exception as e:
                        logger.warning(f"通过注册表查找Chrome失败: {str(e)}")
                
                service = Service(executable_path=chromedriver_path)
                driver = webdriver.Chrome(service=service, options=options)
                
            # 其他浏览器类型的实现可以在这里扩展
            # elif browser_type == BrowserType.FIREFOX:
            #     ...
            
            else:
                raise ValueError(f"不支持的浏览器类型: {browser_type}")
            
            self.active_drivers[driver_id] = driver
            logger.info(f"成功创建浏览器驱动实例: {driver_id}, 浏览器类型: {browser_type}")
            return driver_id
        
        except Exception as e:
            logger.error(f"创建浏览器驱动失败: {str(e)}")
            raise
    
    def get_driver(self, driver_id: str) -> Optional[WebDriver]:
        """获取驱动实例"""
        return self.active_drivers.get(driver_id)
    
    def quit_driver(self, driver_id: str) -> bool:
        """
        退出并关闭驱动实例
        :param driver_id: 驱动ID
        :return: 是否成功关闭
        """
        if driver_id in self.active_drivers:
            try:
                driver = self.active_drivers[driver_id]
                driver.quit()
                del self.active_drivers[driver_id]
                logger.info(f"已关闭浏览器驱动实例: {driver_id}")
                return True
            except Exception as e:
                logger.error(f"关闭浏览器驱动失败: {str(e)}")
                return False
        return False
    
    def quit_all_drivers(self):
        """关闭所有活跃的驱动实例"""
        for driver_id in list(self.active_drivers.keys()):
            self.quit_driver(driver_id)
        logger.info("已关闭所有浏览器驱动实例")
    
    def take_screenshot(self, driver_id: str, filename: str) -> Optional[str]:
        """
        对指定驱动的当前页面进行截图
        :param driver_id: 驱动ID
        :param filename: 文件名
        :return: 截图文件路径或None
        """
        driver = self.get_driver(driver_id)
        if not driver:
            logger.error(f"无法找到驱动实例: {driver_id}")
            return None
        
        try:
            # 确保screenshots目录存在
            import os
            screenshot_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), "..", "screenshots")
            os.makedirs(screenshot_dir, exist_ok=True)
            
            # 生成完整路径
            timestamp = time.strftime("%Y%m%d_%H%M%S")
            filename = f"{filename}_{timestamp}.png"
            filepath = os.path.join(screenshot_dir, filename)
            
            # 截图
            success = driver.save_screenshot(filepath)
            if success:
                logger.info(f"截图成功: {filepath}")
                return filepath
            else:
                logger.error(f"截图失败，未知错误")
                return None
        
        except Exception as e:
            logger.error(f"截图失败: {str(e)}")
            return None

# 单例模式
driver_manager = DriverManager()